package newcrunch.gui.tree;

import java.util.ArrayList;
import java.util.List;
import newcrunch.lexers.JavaTreeLexer;
import newcrunch.lexers.JavaTreeParser;
import org.antlr.runtime.ANTLRStringStream;
import org.antlr.runtime.TokenRewriteStream;
import org.antlr.runtime.tree.CommonTree;
import org.antlr.runtime.tree.Tree;

/**
 * @author WL
 */
public class TreeBuilder 
{    
    private String formatTypeList(Tree node)
    {
        Tree firstNode = node.getChild(0);        
        StringBuilder output = new StringBuilder(formatType(firstNode));
        
        for (int i = 1; i < node.getChildCount(); i++)
        {
            output.append(", ");
            Tree currNode = node.getChild(i);
            output.append(formatType(currNode));
        }
        
        return output.toString();
    }
    
    private CustomTreeNode generateInterfaceNode(Tree node)
    {
        StringBuilder name = new StringBuilder(node.getChild(0).getText());
        StringBuilder toolTip = new StringBuilder("interface " + node.getChild(0).getText());
        CustomTreeNode interfaceNode = new CustomTreeNode(null, TreeNodeType.INTERFACE, null, node);
        
        for (int i = 1; i < node.getChildCount(); i++)
        {
            Tree currNode = node.getChild(i);
            String text = currNode.getText();
            
            if ("GENERICS".equals(text))
            {
                String generics = formatTypeParameter(currNode);
                name.append(generics);
                toolTip.append(generics);
            }
            else if ("EXTENDS".equals(text))
            {
                String superClass = formatType(currNode.getChild(0));
                name.append(" :: " + superClass);
                toolTip.append(" extends " + superClass);
                for (int j = 1; j < currNode.getChildCount(); j++)
                {
                    String innerText = formatType(currNode.getChild(j));
                    name.append(", " + innerText);
                    toolTip.append(", " + innerText);
                }
            }
            else
                addMembers(interfaceNode, currNode);
        }
        
        interfaceNode.setToolTip(toolTip.toString());
        interfaceNode.setUserObject(name.toString());
        
        return interfaceNode;
    }
    
    private String formatTypeArgument(Tree node)
    {
        StringBuilder output = new StringBuilder("<");
        
        for (int i = 0; i < node.getChildCount(); i++)
        {
            String text = node.getChild(i).getText();
            
            if (i > 0)
            {
                if (!node.getChild(i-1).getText().equals("super") &&
                    !node.getChild(i-1).getText().equals("extends"))
                    output.append(",");
                output.append(" ");
            }
            
            if (text.equals("?"))
            {
                output.append("?");
                if (i < node.getChildCount() - 1 && 
                        (node.getChild(i+1).getText().equals("extends") ||
                          node.getChild(i+1).getText().equals("super")))
                {
                    output.append(" " + node.getChild(i+1).getText());
                    i++;
                }
            }
            else
                output.append(formatType(node.getChild(i)));
        }
        
        output.append(">");
        return output.toString();
    }
    
    private String formatType(Tree node)
    {
        if (node.getText().equals("void"))
            return "void";
        
        int j;
        StringBuilder output = new StringBuilder();
        
        for (int i = 0; i < node.getChildCount(); i++)
        {
            Tree currNode = node.getChild(i);
            String text = currNode.getText();
            if (".".equals(text))
            {
                if (!node.getChild(i-1).getText().equals(">") &&
                        !node.getChild(i-1).getText().equals("TYPE_ARG"))
                    output = new StringBuilder();
                else
                    output.append(text);
            }
            else if ("TYPE_ARG".equals(text))
                output.append(formatTypeArgument(currNode));
            else
                output.append(text);
        }
        
        return output.toString();
    }
    
    private CustomTreeNode generateClassNode(Tree tree)
    {
        StringBuilder name = new StringBuilder(tree.getChild(0).getText());
        StringBuilder toolTip = new StringBuilder("class " + tree.getChild(0).getText());
        CustomTreeNode classNode = new CustomTreeNode(null, TreeNodeType.CLASS, null, tree);
        boolean subClassed = false;
        
        for (int i = 1; i < tree.getChildCount(); i++)
        {
            Tree currNode = tree.getChild(i);
            String text = currNode.getText();
            if ("GENERICS".equals(text))
            {
                String generics = formatTypeParameter(currNode);
                name.append(generics);
                toolTip.append(generics);
            }
            else if ("EXTENDS".equals(text))
            {
                subClassed = true;
                //can only extend a single class
                String superClass = formatType(currNode.getChild(0));
                name.append(" :: " + superClass);
                toolTip.append(" extends " + superClass);  
            }
            else if ("IMPLEMENTS".equals(text))
            {
                if (subClassed)
                    name.append(" : ");
                else
                    name.append(" :: ");
                String interfaces = formatTypeList(currNode);
                name.append(interfaces);
                toolTip.append(" implements " + interfaces);
            }
            else if ("CLASSBODY".equals(text))
                for (int j = 0; j < currNode.getChildCount(); j++)
                        addMembers(classNode, currNode.getChild(j));
            else
                toolTip = new StringBuilder(text + " " + toolTip.toString());
        }
                
        classNode.setUserObject(name.toString());
        classNode.setToolTip(toolTip.toString());
        
        return classNode;
    }
    
    private CustomTreeNode generateEnumNode(Tree tree)
    {
        StringBuilder name = new StringBuilder(tree.getChild(0).getText());
        CustomTreeNode enumNode = new CustomTreeNode(null, TreeNodeType.ENUM, null, tree);
        String toolTip = null;
        
        for (int i = 1; i < tree.getChildCount(); i++)
        {
            Tree currNode = tree.getChild(i);
            String currNodeType = currNode.getText();
            
            if ("IMPLEMENTS".equals(currNodeType))
            {
                String interfaces = formatTypeList(currNode).toString();
                name.append(" :: " + interfaces);
                toolTip = "enum " + tree.getChild(0).getText() + " implements " + 
                        interfaces;                
            }
            else if ("ENUM_CONST".equals(currNodeType))
            {
                for (int j = 0; j < currNode.getChildCount(); j++)
                {
                    String enumName = currNode.getChild(j).getText();
                    enumNode.add(new CustomTreeNode(enumName, TreeNodeType.ENUM_CONST, enumName, currNode));
                }                
            }
            else if ("ENUM_DECL".equals(currNodeType))
                for (int j = 0; j < currNode.getChildCount(); j++)
                        addMembers(enumNode, currNode.getChild(j));
        }
        
        enumNode.setUserObject(name.toString());
        enumNode.setToolTip(toolTip == null ? "enum " + name.toString() : toolTip.toString());        
        return enumNode;
    }
    private CustomTreeNode generateAnnotationNode(Tree tree)
    {
        StringBuilder name = new StringBuilder();
        StringBuilder toolTip = new StringBuilder();
        CustomTreeNode annotationNode = new CustomTreeNode(null, TreeNodeType.ANNOTATION, null, tree);
        
        for (int i = 0; i < tree.getChildCount(); i++)
        {
            String text = tree.getChild(i).getText();
            Tree currNode = tree.getChild(i);
            
            if ("MODIFIER".equals(text))
            {
                    toolTip.append(currNode.getChild(0).getText());
                    for (int j = 1; j < currNode.getChildCount(); j++)
                        toolTip.append(" " + currNode.getChild(j).getText());                
            }
            else if ("NAME".equals(text))
            {
                    name.append(currNode.getChild(0).getText());
                    toolTip.append(" @interface " + currNode.getChild(0).getText());                
            }
            else
                addMembers(annotationNode, currNode);
        }
        
        annotationNode.setUserObject(name.toString().trim());
        annotationNode.setToolTip(toolTip.toString().trim());
        
        return annotationNode;
    }
    
    private String formatTypeBound(Tree node)
    {
        StringBuilder result = new StringBuilder("extends");
        
        for (int i = 0; i < node.getChildCount(); i++)
        {
            Tree currNode = node.getChild(i);
            String text = currNode.getText();
            result.append(" ");
            
            if (text.equals("TYPE"))
                result.append(formatType(currNode));
            else
                result.append(text);
        }
        
        return result.toString();
    }
    
    private String formatTypeParameter(Tree node)
    {
        StringBuilder result = new StringBuilder("<");
        
        if (node.getChild(0).getText().equals("TYPE"))
            result.append(formatType(node.getChild(0)));
        else
            result.append(node.getChild(0).getText());
        
        for (int i = 1; i < node.getChildCount(); i++)
        {            
            String text = node.getChild(i).getText();
            Tree currNode = node.getChild(i);
            
            if (text.equals("extends"))
                result.append(" " + formatTypeBound(currNode));
            
            else
                result.append(", " + text);
        }
            
        
        result.append(">");
        return result.toString();
    }
    
    private String formatNormalParameter(Tree node)
    {
        StringBuilder result = new StringBuilder();
        
        //annotations
        int i = 0;
        while (i < node.getChildCount() && !node.getChild(i).getText().equals("TYPE"))
            result.append(node.getChild(i++).getText());
        
        //return type
        if (i < node.getChildCount() && node.getChild(i).getText().equals("TYPE"))
            result.append(" " + formatType(node.getChild(i++)));
        
        //square brackets, if they exist
        while (i < node.getChildCount() && 
                (node.getChild(i).getText().equals("[") || 
                    node.getChild(i).getText().equals("]")))
            result.append(node.getChild(i++));
        
        //variable name
        if (i < node.getChildCount())
            result.append(" " + node.getChild(i++));
        
        return result.toString().trim();
    }
    
    private String formatEllipsisParameter(Tree node)
    {
        StringBuilder result = new StringBuilder();
        
        for (int i = 0; i < node.getChildCount(); i++)
        {
            String text = node.getChild(i).getText();
            if ("final".equals(text))
                result.append("final ");
            else if ("...".equals(text))
                result.append(text);
            else if ("TYPE".equals(text))
                result.append(formatType(node.getChild(i)));
            else
                result.append(" " + text);
        }
        
        return result.toString();
    }
    
    private CustomTreeNode generateMethodNode(Tree node)
    {
        CustomTreeNode result = new CustomTreeNode(null, TreeNodeType.PACKAGE_METHOD, null, node);
        StringBuilder toolTip = new StringBuilder();
        StringBuilder displayBuf = new StringBuilder();
        
        String returnDisplay = "";
        String name = "";
        
        for (int i = 0; i < node.getChildCount(); i++)
        {
            Tree currNode = node.getChild(i);
            String nodeType = currNode.getText();
            
            if ("MODIFIER".equals(nodeType))
            {
                for (int j = 0; j < currNode.getChildCount(); j++)
                {
                    String innerNodeType = currNode.getChild(j).getText();
                    toolTip.append(" " + innerNodeType);
                    if ("public".equals(innerNodeType))
                        result.setType(TreeNodeType.PUBLIC_METHOD);
                    else if ("private".equals(innerNodeType))
                        result.setType(TreeNodeType.PRIVATE_METHOD);
                    else if ("protected".equals(innerNodeType))
                        result.setType(TreeNodeType.PROTECTED_METHOD);
                }                
            }
            else if ("TYPE".equals(nodeType))
                toolTip.append(" " + formatTypeParameter(currNode));
            else if ("RETURN".equals(nodeType))
            {
                returnDisplay = formatType(currNode.getChild(0));
                toolTip.append(" " + returnDisplay);
            }
            else if ("NAME".equals(nodeType))
            {
                    name = currNode.getChild(0).getText();
                    displayBuf.append(name);                
            }
            else if ("PARA".equals(nodeType))
            {
                StringBuilder parameters = new StringBuilder();
                boolean appended = false;
                //ignore the "(" and ")"
                for (int j = 0; j < currNode.getChildCount(); j++)
                {
                    Tree innerNode = currNode.getChild(j);
                    String innerNodeText = innerNode.getText();
                    if ("NORMAL_PARA".equals(innerNodeText))
                    {
                        if (!appended)
                            appended = true;
                        else
                        {
                            parameters.append(", ");
                            appended = true;
                        }
                        parameters.append(formatNormalParameter(innerNode));
                    }
                    else if ("ELLIPSIS_PARA".equals(innerNodeText))
                    {
                        if (!appended)
                            appended = true;
                        else
                        {
                            parameters.append(", ");
                            appended = true;
                        }
                        parameters.append(formatEllipsisParameter(innerNode));
                    }
                    else if ("[".equals(innerNodeText) || "]".equals(innerNodeText))
                    {
                        returnDisplay += innerNodeText;
                        toolTip.append(innerNodeText);
                    }
                    else if ("(".equals(innerNodeText) || ")".equals(innerNodeText))
                    {
                    }
                    else
                        parameters.append(innerNodeText);
                }

                String parameterDisplay = "(" + parameters.toString() + ")";
                toolTip.append(" " + name + parameterDisplay);
                displayBuf.append(parameterDisplay);                
            }
            else if ("THROWS".equals(nodeType))
            {
                toolTip.append(" throws ");
                for (int j = 0; j < currNode.getChildCount(); j++)
                {
                    if (j > 0)
                        toolTip.append(", ");

                    toolTip.append(currNode.getChild(j).getChild(currNode.getChild(j).getChildCount()-1));
                }                
            }
        }
        if (!returnDisplay.equals(""))
            displayBuf.append(" : " + returnDisplay);
        result.setToolTip(toolTip.toString().trim());
        result.setUserObject(displayBuf.toString());
        return result;
    }
    
    private List<CustomTreeNode> generateFieldNode(Tree node)
    {
        List<CustomTreeNode> fieldNodes = new ArrayList<CustomTreeNode>();
        StringBuilder modifier = new StringBuilder();
        String type = "";
        TreeNodeType currType = TreeNodeType.PACKAGE_FIELD;
        
        for (int i = 0; i < node.getChildCount(); i++)
        {
            String currText = node.getChild(i).getText();
            
            if ("MODIFIER".equals(currText))
                for (int j = 0; j < node.getChild(i).getChildCount(); j++)
                    {
                        String text = node.getChild(i).getChild(j).getText();
                        if ("private".equals(text))
                            currType = TreeNodeType.PRIVATE_FIELD;
                        
                        else if ("protected".equals(text))
                            currType = TreeNodeType.PROTECTED_FIELD;
                        
                        else if ("public".equals(text))
                            currType = TreeNodeType.PUBLIC_FIELD;
                        modifier.append(text + " ");
                    }
            else if ("TYPE".equals(currText))
                type = formatType(node.getChild(i));
            
            else if ("VARIABLE".equals(currText))
            {
                //for the case of int[] x[], y
                //type of x will be int[][]
                //whereas y is int[]
                String tempType = "";
                for (int j = 0; j < node.getChild(i).getChildCount(); j++)
                {
                    String text = node.getChild(i).getChild(j).getText();

                    if ("[".equals(text) || "]".equals(text))
                        tempType += text;
                    else
                    {
                        CustomTreeNode fieldNode = new CustomTreeNode(null, currType, null, node.getChild(i).getChild(j));

                        fieldNode.setUserObject(text + " : " + type + tempType);
                        fieldNode.setType(currType);
                        fieldNode.setToolTip((modifier.toString().trim() + " " + type + tempType + " " + text).trim());

                        fieldNodes.add(fieldNode);
                        tempType = "";
                    }
                }
            }
        }
        
        return fieldNodes;
    }
    
    private CustomTreeNode generateAnnotationMethod(Tree node)
    {
        //modifier can only be public
        CustomTreeNode result = new CustomTreeNode(null, TreeNodeType.PUBLIC_METHOD, null, node);
        
        String type = formatType(node.getChild(0));
        String name = node.getChild(1).getText() + " : " + type;
        String toolTip = "public abstract " + type + " " + node.getChild(1).getText();
        
        result.setToolTip(toolTip);
        result.setUserObject(name);
        
        return result;
    }
    
    /*
     * Add nodes originating from source to root.
     */
    private void addMembers(CustomTreeNode root, Tree source)
    {
        String nodeType = source.getText();
        
        if ("CLASS".equals(nodeType))
            root.add(generateClassNode(source));
        
        else if ("INTERFACE".equals(nodeType))
            root.add(generateInterfaceNode(source));
        
        else if ("ENUM".equals(nodeType))
            root.add(generateEnumNode(source));

        else if ("ANNOTATION".equals(nodeType))
            root.add(generateAnnotationNode(source));
        
        else if ("METHOD".equals(nodeType))
            root.add(generateMethodNode(source));
        
        else if ("FIELD".equals(nodeType))
        {
            List<CustomTreeNode> result = generateFieldNode(source);
            for (CustomTreeNode node : result)
                root.add(node);
        }
        
        else if ("ANNOTATION_METHOD".equals(nodeType))
            root.add(generateAnnotationMethod(source));
    }
    
    /**
     * Regenerates the entire tree
     */
    public CustomTreeNode generateTree(String input)
    {
        try {            
            //parse the code and generate the AST
            ANTLRStringStream ss = new ANTLRStringStream(input);
            //ANTLRStringStream ss = new ANTLRStringStream(input);
            JavaTreeLexer lex = new JavaTreeLexer(ss);
            TokenRewriteStream tokens = new TokenRewriteStream(lex);
            JavaTreeParser grammar = new JavaTreeParser(tokens);
            
            //obtaining the root node from the AST
            JavaTreeParser.compilationUnit_return ret = grammar.compilationUnit();
            Tree root = (Tree)ret.getTree();
                        
            //generate the JTree nodes using the AST
            CustomTreeNode rootNode = new CustomTreeNode("ROOT", TreeNodeType.CLASS, null, null);
            CommonTree currNode = null;
            String currNodeType = null;
            
            //iterate through the top level nodes
            for (int i = 0; i < root.getChildCount(); i++)
            {
                currNode = (CommonTree)root.getChild(i);
                
                if (currNode == null || currNode.getToken() == null)
                    continue;
                
                currNodeType = currNode.getToken().getText();
                
                if ("CLASS".equals(currNodeType))
                    rootNode.add(generateClassNode(currNode));
                
                else if ("INTERFACE".equals(currNodeType))
                    rootNode.add(generateInterfaceNode(currNode));
                
                else if ("ENUM".equals(currNodeType))
                    rootNode.add(generateEnumNode(currNode));
                
                else if ("ANNOTATION".equals(currNodeType))
                    rootNode.add(generateAnnotationNode(currNode));
            }
            
            //recursively sorts the root
            new TreeSorter().sort(rootNode);
                        
            return rootNode;
        } catch (org.antlr.runtime.RecognitionException e) {
            e.printStackTrace();
        }
        
        return null;
    }
}
